#ifndef SSY_MATRIX_DIAGONAL_H
#define SSY_MATRIX_DIAGONAL_H


#include "MutrixBase.h"

template<typename ElementType>
class MatrixDiagonal : public MatrixBase<ElementType> {
public:
	MatrixDiagonal(size_t series) : MatrixBase(series, series) {
		SetType(MATRIX_DIAGONAL);
		_mut.resize(GetRow());
	}
	explicit MatrixDiagonal(const MatrixBase& matbase);
	~MatrixDiagonal(void) {
		vector<ElementType>().swap(_mut);
	}

	const ElementType& GetValue(size_t rowidx, size_t colidx) const;
	bool SetValue(size_t rowidx, size_t colidx, const ElementType& element);
	void Print(void) const;

	MatrixDiagonal<ElementType> Transpose(void) const { return (*this); }
	MatrixDiagonal<ElementType>& operator=(const MatrixBase<ElementType>& matbase);
	const ElementType& operator[](size_t row) const { return _mut[row]; } 
	ElementType& operator[](size_t row) { 
		return const_cast<ElementType&>
			(static_cast<const MatrixDiagonal<ElementType>&>(*this)[row]); 
	}

	friend const MatrixDiagonal operator+(const MatrixDiagonal& matdiag1, const MatrixDiagonal& matdiag2) {
		return doAdd(matdiag1, matdiag2);
	}
	friend const MatrixDiagonal operator-(const MatrixDiagonal& matdiag1, const MatrixDiagonal& matdiag2) {
		return doSubtract(matdiag1, matdiag2);
	}
	friend const MatrixDiagonal operator*(const MatrixDiagonal& matdiag1, const MatrixDiagonal& matdiag2) {
		return doMultiply(matdiag1, matdiag2);
	}
	friend const MatrixDiagonal operator*(const MatrixDiagonal& matdiag, const ElementType& element) {
		return doElementMultiply(matdiag, element);
	}
	friend const MatrixDiagonal operator*(const ElementType& element, const MatrixDiagonal& matdiag) {
		return doElementMultiply(matdiag, element);
	}

private:
	bool PushValue(const ElementType& element);

	vector<ElementType> _mut; 
};

template<typename ElementType>
MatrixDiagonal<ElementType>::MatrixDiagonal(const MatrixBase& matbase) 
	: MatrixBase(matbase) {
	_mut.resize(GetRow());
	for (size_t rIdx = 0; rIdx < GetRow(); rIdx++) {
		for (size_t cIdx = 0; cIdx < GetCol(); cIdx++) {
			if (rIdx == cIdx)
				assert(SetValue(rIdx, cIdx, matbase.GetValue(rIdx, cIdx)));
			else
				assert(_c_zero == matbase.GetValue(rIdx, cIdx));
		}
	}
}

template<typename ElementType>
bool MatrixDiagonal<ElementType>::PushValue(const ElementType& element) {
	if (GetRow() != GetCol()) {
		SetErrMsg("This mutrix is uncorrect, row != col");
		if (_rowidx == GetRow()-1 && _colidx == GetCol()) 
			SetErrMsg("This mutrix is full, can't push");
		return false;
	}
	else {
		_mut[_rowidx] = element;
		_rowidx++;
		return true;
	}
}

template<typename ElementType>
const ElementType& MatrixDiagonal<ElementType>::GetValue(size_t rowidx, size_t colidx) const {
	assert(rowidx < GetRow());
	assert(colidx < GetCol());
	if (rowidx != colidx) 
		return _c_zero;
	else
		return _mut[rowidx];
}

template<typename ElementType>
bool MatrixDiagonal<ElementType>::SetValue(size_t rowidx, size_t colidx, const ElementType& element) {
	assert(rowidx < GetRow());
	assert(colidx < GetCol());
	if (rowidx != colidx) {
		if (element != _c_zero) {
			SetErrMsg("SetValue() input parameter is uncorrect");
			return false;
		}
	} 
	else
		_mut[rowidx] = element;
	return true;
}

template<typename ElementType>
void MatrixDiagonal<ElementType>::Print(void) const {
	cout << "[" << GetRow() << "," << GetCol() << "] "; 
	for (size_t rIdx = 0; rIdx < GetRow(); rIdx++) {
		for (size_t cIdx = 0; cIdx < GetCol(); cIdx++) {
			if (rIdx == cIdx)
				cout << _mut[rIdx] << " ";
			else
				cout << _c_zero << " ";
		}
	}
	cout << endl;
}

template<typename ElementType>
MatrixDiagonal<ElementType>& MatrixDiagonal<ElementType>::operator=(const MatrixBase<ElementType>& matbase) {
	if (this != &matbase)
		MatrixBase::operator=(matbase);	
	return (*this);
}

template<typename ElementType>
const MatrixDiagonal<ElementType> doAdd(const MatrixDiagonal<ElementType>& matdiag1, 
										const MatrixDiagonal<ElementType>& matdiag2) {
	assert(matdiag1.GetRow() == matdiag2.GetRow());
	assert(matdiag1.GetCol() == matdiag2.GetCol());
	assert(matdiag1.GetRow() == matdiag1.GetCol());
	MatrixDiagonal<ElementType> matsum(matdiag1);
	ElementType element;
	for (size_t rIdx = 0; rIdx < matsum.GetRow(); rIdx++) {
		element = matdiag1.GetValue(rIdx, rIdx) + matdiag2.GetValue(rIdx, rIdx);
		matsum.SetValue(rIdx, rIdx, element);
	}
	return matsum;
}

template<typename ElementType>
const MatrixDiagonal<ElementType> doSubtract(const MatrixDiagonal<ElementType>& matdiag1, 
										const MatrixDiagonal<ElementType>& matdiag2) {
	assert(matdiag1.GetRow() == matdiag2.GetRow());
	assert(matdiag1.GetCol() == matdiag2.GetCol());
	assert(matdiag1.GetRow() == matdiag1.GetCol());
	MatrixDiagonal<ElementType> matdiff(matdiag1);
	ElementType element;
	for (size_t rIdx = 0; rIdx < matdiff.GetRow(); rIdx++) {
		element = matdiag1.GetValue(rIdx, rIdx) - matdiag2.GetValue(rIdx, rIdx);
		matdiff.SetValue(rIdx, rIdx, element);
	}
	return matdiff;
}

template<typename ElementType>
const MatrixDiagonal<ElementType> doMultiply(const MatrixDiagonal<ElementType>& matdiag1, 
									const MatrixDiagonal<ElementType>& matdiag2) {
	assert(matdiag1.GetRow() == matdiag2.GetCol());
	assert(matdiag1.GetCol() == matdiag2.GetRow());
	assert(matdiag1.GetRow() == matdiag1.GetCol());
	MatrixDiagonal<ElementType> matproduct(matdiag1.GetRow(), matdiag2.GetCol());
	ElementType element;
	for (size_t rIdx = 0; rIdx < matdiag1.GetRow(); rIdx++) {
		element = matdiag1.GetValue(rIdx, rIdx) * matdiag2.GetValue(rIdx, rIdx);
		assert(matproduct.SetValue(rIdx, rIdx, element));
	}
	return matproduct;
}

template<typename ElementType>
const MatrixDiagonal<ElementType> doElementMultiply(const MatrixDiagonal<ElementType>& matdiag, 
												  const ElementType& element) {
	MatrixDiagonal<ElementType> matproduct(matdiag);
	for (size_t rIdx = 0; rIdx < matproduct.GetRow(); rIdx++) {
		assert(matproduct.SetValue(rIdx, rIdx, matproduct.GetValue(rIdx, rIdx) * element));
	}
	return matproduct;
}


#endif // SSY_MATRIX_DIAGONAL_H


